import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
import pandas as pd


def plot_semantic_model():
    data = {'Model': ['BGE', 'BGE', 'BGE-m3', 'BGE-m3', 'BGE-reranker', 'BGE-reranker', 'LLM-embedder', 'LLM-embedder',
                      'SIMCSE', 'SIMCSE'],
            'Precision': [64.78, 61.50, 61.99, 58.92, 28.08, 20.34, 41.59, 31.58, 60.98, 56.71],
            'Group': ["macro", "micro", "macro", "micro", "macro", "micro", "macro", "micro", "macro", "micro"]}
    df = pd.DataFrame(data)
    sns.barplot(x='Model', y='Precision', hue='Group', data=df)
    plt.xlabel('Models')
    plt.ylabel('Precision')
    # plt.title('Bar Chart using Seaborn')
    plt.savefig('semantic_model.eps', format='eps')
    plt.show()


def plot_model_scale():
    data = {'Scale': ['Small', 'Base', 'large'],
            'Original-macro': [55.65, 57.66, 58.62],
            'Original-micro': [51.12, 53.68, 55.00],
            'Finetuned-macro': [59.12, 63.09, 64.78],
            'Finetuned-micro': [54.66, 59.00, 61.50]}
    bar_width1 = 0.5  # 初始值柱子宽度
    bar_width2 = 0.5  # 增长后值柱子宽度
    index = np.arange(len(data['Scale']))

    # 绘制初始值柱子，向右偏移一点
    plt.bar(index*2 - 0.1, data['Original-macro'], bar_width1, label='Original', color='skyblue', zorder=2)

    # 绘制增长后值柱子，向左偏移一点
    plt.bar(index*2 + 0.1, data["Finetuned-macro"], bar_width2, label='Classical Adapted', color='lightgreen', zorder=1)

    plt.bar(index*2 + 0.8 - 0.1, data['Original-micro'], bar_width1,  color='skyblue', zorder=2)

    # 绘制增长后值柱子，向左偏移一点
    plt.bar(index*2 + 0.8 + 0.1, data["Finetuned-micro"], bar_width2,  color='lightgreen', zorder=1)
    # 绘制线段表示增长
    for i in range(len(data['Scale'])):
        plt.plot([index[i]*2 - 0.1, index[i]*2 + 0.1], [data['Original-macro'][i], data['Finetuned-macro'][i]], color='blue', linestyle='--',
                 marker='o')
        growth = data['Finetuned-macro'][i] - data['Original-macro'][i]
        plt.text(index[i]*2, data['Finetuned-macro'][i] + 0.5, f'+%.1f'%growth, ha='center', va='bottom')

        plt.text(index[i]*2, data['Finetuned-macro'][i] + 5, "macro", ha='center', va='top')
        plt.plot([index[i]*2 + 0.8 - 0.1, index[i]*2 +0.8 + 0.1], [data['Original-micro'][i], data['Finetuned-micro'][i]], color='blue', linestyle='--',
                 marker='o')
        growth = data['Finetuned-micro'][i] - data['Original-micro'][i]
        plt.text(index[i]*2+0.6, data['Finetuned-micro'][i] + 0.5, f'+%.1f'%growth, ha='center', va='bottom')
        plt.text(index[i]*2+0.7, data['Finetuned-micro'][i] + 5, "micro", ha='center', va='top')

    plt.xlabel('Model Scale')
    plt.ylabel('Precision')
    plt.ylim([30, 80])
    plt.xticks([2*i+0.3 for i in index], data['Scale'])
    plt.legend()

    plt.savefig('model_scale.eps', format='eps')
    plt.show()


def plot_community_sim():
    data = {'Semantic': ['BGE', 'BGE scenario', 'Hybrid'],
            'macro w/o. cluster-level similarity': [76.58, 81.12, 81.79],
            'micro w/o. cluster-level similarity': [76.48, 83.75, 85.93],
            'macro w. cluster-level similarity': [76.85, 81.59, 82.17],
            'micro w. cluster-level similarity': [77.11, 84.81, 86.55]}
    bar_width1 = 0.5  # 初始值柱子宽度
    bar_width2 = 0.5  # 增长后值柱子宽度
    index = np.arange(len(data['Semantic']))

    # 绘制初始值柱子，向右偏移一点
    plt.bar(index*2 - 0.1, data['macro w/o. cluster-level similarity'], bar_width1, label='w/o. Cluster Sim', color='crimson', zorder=2)

    # 绘制增长后值柱子，向左偏移一点
    plt.bar(index*2 + 0.1, data["macro w. cluster-level similarity"], bar_width2, label='w. Cluster Sim', color='forestgreen', zorder=1)

    plt.bar(index*2 + 0.8 - 0.1, data['micro w/o. cluster-level similarity'], bar_width1,  color='crimson', zorder=2)

    # 绘制增长后值柱子，向左偏移一点
    plt.bar(index*2 + 0.8 + 0.1, data["micro w. cluster-level similarity"], bar_width2,  color='forestgreen', zorder=1)
    # 绘制线段表示增长
    for i in range(len(data['Semantic'])):
        plt.plot([index[i]*2 - 0.1, index[i]*2 + 0.1], [data['macro w/o. cluster-level similarity'][i], data['macro w. cluster-level similarity'][i]], color='blue', linestyle='--',
                 marker='o')
        growth = data['macro w. cluster-level similarity'][i] - data['macro w/o. cluster-level similarity'][i]
        plt.text(index[i]*2, data['macro w. cluster-level similarity'][i] + 0.3, f'+%.1f'%growth, ha='center', va='bottom')

        plt.text(index[i]*2, data['macro w. cluster-level similarity'][i] + 1.5, "macro", ha='center', va='top')
        plt.plot([index[i]*2 + 0.8 - 0.1, index[i]*2 +0.8 + 0.1], [data['micro w/o. cluster-level similarity'][i], data['micro w. cluster-level similarity'][i]], color='blue', linestyle='--',
                 marker='o')
        growth = data['micro w. cluster-level similarity'][i] - data['micro w/o. cluster-level similarity'][i]
        plt.text(index[i]*2+0.6, data['micro w. cluster-level similarity'][i] + 0.3, f'+%.1f'%growth, ha='center', va='bottom')
        plt.text(index[i]*2+0.7, data['micro w. cluster-level similarity'][i] + 1.5, "micro", ha='center', va='top')

    plt.xlabel('Semantic Encoding')
    plt.ylabel('Precision')
    plt.ylim([75, 92])
    plt.xticks([2*i+0.3 for i in index], data['Semantic'])
    plt.legend()

    plt.savefig('community_sim.eps', format='eps')
    plt.show()

def plot_dp():
    data = {'macro': [64.78, 49.79, 81.12],
            'micro': [61.50, 31.54, 83.75],
            'strategy': ['greedy', 'w/o. dp', 'w. dp'],}
    bar_width1 = 0.8  # 初始值柱子宽度
    bar_width2 = 0.8  # 增长后值柱子宽度
    gap = 0.4
    index = np.arange(len(data['strategy']))
    double_index = np.concatenate([index+gap/2, index+gap + len(index)], axis=0)
    # 绘制初始值柱子，向右偏移一点
    plt.bar(index+gap/2, data['macro'], bar_width1, label='macro precision', color='#FF6347', zorder=1)

    # 绘制增长后值柱子，向左偏移一点
    plt.bar([i+len(index)+gap for i in index], data["micro"], bar_width2, label='micro precision', color='#4682B4', zorder=1)

    # 绘制线段表示增长
    plt.xlabel('Optimization Strategy')
    plt.ylabel('Precision')
    plt.ylim([20, 90])
    plt.xticks(double_index, data['strategy']*2)
    plt.legend()

    plt.savefig('optimization_strategy.eps', format='eps')
    plt.show()


if __name__ == '__main__':
    plot_community_sim()
